import numpy as np
import os
import pandas as pd
import filter_unused_data


def convert_ratings(input_fname, output_fname):
    rating = pd.DataFrame([], columns=['userID', 'itemID'])
    n_file = 16
    for fileID in range(n_file):
        f = np.load(os.path.join(input_fname, f'testx16x32_{fileID}.npz'))
        # print(list(f.keys()))
        interact_l = np.array(f['arr_0'])
        df2 = pd.DataFrame(interact_l, columns=['userID', 'itemID'])
        rating = pd.concat([rating, df2])
        # print(f['arr_0'])
    rating['rating'] = 1

    # print(rating)

    rating = filter_unused_data.convert_ratings(rating)

    filter_unused_data.stats_rating(rating)

    order_l = ['userID', 'itemID', 'rating']
    rating = rating[order_l]
    rating.to_csv(output_fname, index=False)
    print("complete {}".format(output_fname))


if __name__ == '__main__':
    basic_dir = '/home/bianzheng/Dataset/Recommendation/movielens-1b/ml-20mx16x32'
    convert_ratings(input_fname=basic_dir,
                    output_fname='/home/bianzheng/rec2-mips/intermediate-rating-csv/movielens-1b.csv')
